from knowledge_base.generate import generate

from knowledge_base.query_rewrite import query_rewrite
from knowledge_base.recall import recall
from knowledge_base.rerank import rerank


class KnowledgeBasePipeline:
    def __init__(self, query, context, filter_files, permission_list, q_type):
        self.query = query
        self.context = context
        self.filter_files = filter_files
        self.permission_list = permission_list
        self.q_type = q_type

    async def process(self, stream=False):
        rewrite_result = await query_rewrite(self.query, self.context)
        recall_result = recall(rewrite_result, self.filter_files, self.permission_list, self.q_type)
        rerank_result = rerank(recall_result)
        if not stream:
            return await generate(rerank_result, self.context, self.query, self.q_type)
        else:
            return await generate(rerank_result, self.context, self.query, self.q_type, stream=True)
